import cv2
import math
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve2d


def compute_mse(img1, img2):
    mse = np.mean((img1 / 255. - img2 / 255.) ** 2)
    return mse


def compute_psnr(img1, img2):
    mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)
    if mse < 1.0e-10:
        return 100
    return 10 * math.log10(255.0 ** 2 / mse)


def matlab_style_gauss2D(shape=(3, 3), sigma=0.5):
    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m + 1, -n:n + 1]
    h = np.exp(-(x * x + y * y) / (2. * sigma * sigma))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    return h


def filter2(x, kernel, mode='same'):
    return convolve2d(x, np.rot90(kernel, 2), mode=mode)


def compute_ssim(im1, im2, k1=0.01, k2=0.03, win_size=11, L=255):
    if not im1.shape == im2.shape:
        raise ValueError("Input Imagees must have the same dimensions")
    if len(im1.shape) > 2:
        raise ValueError("Please input the images with 1 channel")
    M, N = im1.shape
    C1 = (k1 * L) ** 2
    C2 = (k2 * L) ** 2
    window = matlab_style_gauss2D(shape=(win_size, win_size), sigma=1.5)
    window = window / np.sum(np.sum(window))

    if im1.dtype == np.uint8:
        im1 = np.double(im1)
    if im2.dtype == np.uint8:
        im2 = np.double(im2)

    mu1 = filter2(im1, window, 'valid')
    mu2 = filter2(im2, window, 'valid')
    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = filter2(im1 * im1, window, 'valid') - mu1_sq
    sigma2_sq = filter2(im2 * im2, window, 'valid') - mu2_sq
    sigmal2 = filter2(im1 * im2, window, 'valid') - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigmal2 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
    return np.mean(np.mean(ssim_map))


if __name__ == '__main__':
    img_origin = cv2.imread(r"C:\Users\Public\opencv\Figure\lena.jpg", 0)
    plt.subplot(421)
    plt.imshow(img_origin)
    img_Mb = cv2.imread(r'C:\Users\Public\opencv\Figure\lenaMb.jpg', 0)
    plt.subplot(422)
    plt.imshow(img_Mb)

    img_inv = cv2.imread(r'C:\Users\Public\opencv\Figure\lenaInv.jpg', 0)
    plt.subplot(423)
    plt.imshow(img_inv)
    img_invn = cv2.imread(r'C:\Users\Public\opencv\Figure\lenaInvn.jpg', 0)
    plt.subplot(424)
    plt.imshow(img_invn)
    # 原图与运动模糊后的逆滤波图像之间的评价参数
    mse = compute_mse(img_origin, img_inv)
    print('MSE:{}'.format(mse))
    psnr = compute_psnr(img_origin, img_inv)
    print('PSNR:{}'.format(psnr))
    ssim = compute_ssim(img_origin, img_inv)
    print('SSIM:{}'.format(ssim))
    # 原图与运动模糊+噪声后的逆滤波图像之间的评价参数
    mse = compute_mse(img_origin, img_invn)
    print('MSE:{}'.format(mse))
    psnr = compute_psnr(img_origin, img_invn)
    print('PSNR:{}'.format(psnr))
    ssim = compute_ssim(img_origin, img_invn)
    print('SSIM:{}'.format(ssim))

    img_origin = cv2.imread(r"C:\Users\Public\opencv\Figure\lena.jpg", 0)
    plt.subplot(425)
    plt.imshow(img_origin)

    img_Mb = cv2.imread(r'C:\Users\Public\opencv\Figure\lenaMb.jpg', 0)
    plt.subplot(426)
    plt.imshow(img_Mb)

    img_wd = cv2.imread(r"C:\Users\Public\opencv\Figure\lenaMd.jpg", 0)
    plt.subplot(427)
    plt.imshow(img_wd)

    img_wdn = cv2.imread(r"C:\Users\Public\opencv\Figure\lenaWdn.jpg", 0)
    plt.subplot(428)
    plt.imshow(img_wdn)

    # 原图与运动模糊后的维纳滤波图像之间的评价参数
    mse2 = compute_mse(img_origin, img_wd)
    print('MSE:{}'.format(mse2))
    psnr2 = compute_psnr(img_origin, img_wd)
    print('PSNR:{}'.format(psnr2))
    ssim2 = compute_ssim(img_origin, img_wd)
    print('SSIM:{}'.format(ssim2))
    # 原图与运动模糊+噪声后的维纳滤波图像之间的评价参数
    mse2 = compute_mse(img_origin, img_wdn)
    print('MSE:{}'.format(mse2))
    psnr2 = compute_psnr(img_origin, img_wdn)
    print('PSNR:{}'.format(psnr2))
    ssim2 = compute_ssim(img_origin, img_wdn)
    print('SSIM:{}'.format(ssim2))
    cv2.waitKey(0)
    cv2.destroyAllWindows()